#!/usr/bin/env python
# -*-coding:utf-8-*-



from ner_predict import Predictor
import conf


predictor = Predictor(map_file=conf.NER_MAP_PATH, checkpoint_dir=conf.NER_CHECKPOINT_DIR)


def ner_by_local_model(words, segs):
    ners = predictor.predict(words, segs)

    result = []
    entity = ""
    i = 0

    for label, word in zip(ners, words):
        i += 1
        if label == "O":
            result.append((i, i + 1, label, word))
            continue

        lt = label.split("_")[1]
        lb = label.split("_")[0]

        if lb == "S":
            result.append((i, i + 1, lt, word))

        elif lb == "B":
            entity += word

        elif lb == "M":
            entity += word
        elif lb == "E":
            entity += word
            result.append((i - len(entity), i + 1, lt, entity))
            entity = ""
    return result


if __name__ == '__main__':
    text = "继前不久始于中国的召回风波，宝马因为车辆的发动机螺栓故障，在全球范围将召回48.9万辆车，在原有中国召回的基础上数量进一步增加。据悉，召回车辆将包括北美市场的15.6万辆，宝马曾于3月宣布在华召回232,098辆发动机螺栓故障车辆。涉及车型包括搭载六缸发动机的宝马5系、7系、X3、X5，。但具体型号Santer并没有透露。宝马发言人Bernhard Santer表示，目前尚无该故障造成事故或伤亡的报告。但他仍建议相关车主及时检查车辆引擎。Santer说，凭借剩余的动力，车辆仍旧可以坚持到最近的修理厂。"
    text = "　　那么，如何构建“网络空间命运共同体”？早在2015年，习近平站在乌镇的讲台上发表重要讲话时，就明确指出应坚持“四项原则”，并提出“五点主张”。他提出推动互联网全球治理体系变革，建立多边、民主、透明的全球互联网治理体系，倡议各国应该加强沟通、扩大共识、深化合作，共同构建网络空间命运共同体。此次主旨演讲，一度被海内外称为世界互联网治理的中国方案。"
    words = list(text)
    ners = ner_by_local_model(words)
    for n in ners:
        print(n)
